package batchMallocer

import (
	"fmt"
)

/*
	如果有32个elem,则32/30=1.067,这样每一个片有1个,余下2个,余数必定小于30,
	所以我们将余数平摊到所有的片中,只要在每一个片中多加一个elem即可,
	直到加完不超过elem总数32,

	就这样前两个槽中就每一个多一个,
	1 1 1 1 1 1 1 .... 1 1 1 (30个)
	1 1

	主要在总数为小数量的情况下,这样的分布能达到比较均匀的效果,
	不过在大数量的情况下,因为余数本身小于30,所以相对一个片的数量如果较小的话,
	基本平不平摊也没有太大所谓,不过从理想的角度,这样的分布是最好的
*/
func GenBatchEndIdxes(totalCnt int, batchCnt int) ([]int, int /*final batch cnt*/) {
	baseCntInBatch := totalCnt / batchCnt
	cntInBatchs := make([]int, batchCnt)
	for i := 0; i < batchCnt; i++ {
		cntInBatchs[i] = baseCntInBatch
	}
	leftCnt := totalCnt - baseCntInBatch*batchCnt
	for i := 0; i < leftCnt; i++ {
		cntInBatchs[i] += 1
	}
	if leftCnt == totalCnt {
		batchCnt = totalCnt
	}

	endIdxs := cntInBatchs
	for i := 0; i < batchCnt; i++ {
		if i > 0 {
			endIdxs[i] += endIdxs[i-1]
		}
	}

	return endIdxs[:batchCnt], batchCnt
}

func GetBeginAndEndIdx(idx int, batchEndIdxes []int) (beginIdx int, endIdx int) {
	idx %= len(batchEndIdxes)

	beginIdx = 0
	if idx > 0 {
		beginIdx = batchEndIdxes[idx-1]
	}
	endIdx = batchEndIdxes[idx]
	return
}

// 这里只有beforeProc,没有afterProc,因为proc可能是一个go调用, 执行内容到其他goroutine中去了, 所以即使在函数最后加上afterProc, 也不能正确反映proc的结束
func ProcBatches(totalCnt int, maxBatchCnt int, beforeProc func(finalBatchCnt int), proc func(beginIdx int, endIdx int) bool) error {
	if totalCnt <= 0 {
		return nil
	}
	if maxBatchCnt <= 0 {
		return fmt.Errorf("maxBatchCnt(%d) should be > 0", maxBatchCnt)
	}

	endIdxs, finalBatchCnt := GenBatchEndIdxes(totalCnt, maxBatchCnt)

	if beforeProc != nil {
		beforeProc(finalBatchCnt)
	}

	for i := 0; i < finalBatchCnt; i++ {
		beginIdx, endIdx := GetBeginAndEndIdx(i, endIdxs)
		if beginIdx == endIdx {
			continue
		}

		if !proc(beginIdx, endIdx) {
			return nil
		}
	}
	return nil
}

func ProcBatchesByCnt(totalCnt int, batchUnitCnt int, beforeProc func(finalBatchCnt int), proc func(beginIdx int, endIdx int) bool) error {
	batchCnt := totalCnt / batchUnitCnt
	if batchCnt <= 0 {
		batchCnt = 1
	}
	return ProcBatches(totalCnt, batchCnt, beforeProc, proc)
}

type DynamicBatchMallocer struct {
	endIdxs       []int
	finalBatchCnt int
	curBatchIdx   int
}

func NewDynamicBatchMallocer(totalCnt int, maxBatchCnt int) *DynamicBatchMallocer {
	if maxBatchCnt <= 0 {
		panic(fmt.Errorf("maxBatchCnt(%d) should be > 0", maxBatchCnt))
	}

	endIdxs, finalBatchCnt := GenBatchEndIdxes(totalCnt, maxBatchCnt)
	rlt := &DynamicBatchMallocer{
		endIdxs:       endIdxs,
		finalBatchCnt: finalBatchCnt,
		curBatchIdx:   -1,
	}
	return rlt
}

func (self *DynamicBatchMallocer) GetFinalBatchCnt() int {
	return self.finalBatchCnt
}

func (self *DynamicBatchMallocer) GetNextBatch(loop bool) (beginIdx int, endIdx int, ok bool) {
	if self.finalBatchCnt <= 0 {
		return -1, -1, false
	}

	if !loop && self.curBatchIdx == len(self.endIdxs)-1 {
		return -1, -1, false
	}

	self.curBatchIdx = (self.curBatchIdx + 1) % len(self.endIdxs)
	beginIdx, endIdx = GetBeginAndEndIdx(self.curBatchIdx, self.endIdxs)
	return beginIdx, endIdx, true
}
